package org.example.core;

import org.antlr.v4.runtime.CharStream;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.tree.ParseTree;
import org.example.grammar.SZCGrammarBaseVisitor;
import org.example.grammar.SZCGrammarLexer;
import org.example.grammar.SZCGrammarParser;
import org.example.sym.SymbolTable;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

public class MyVisitor extends SZCGrammarBaseVisitor<String> {
    private final StringBuilder outputIR = new StringBuilder();
    private long tempSymCount = 0;
    private long ifUseCount = 0;
    private long whileUseCount = 0;
    private final SymbolTable symbolTable = new SymbolTable(0, null);

    private String getAUniqueTempSym() {
        String retSym = "%" + tempSymCount;
        tempSymCount++;
        return retSym;
    }

    private String getAIfNum() {
        String retNum = String.valueOf(ifUseCount);
        ifUseCount++;
        return retNum;
    }

    private String getAWhileNum() {
        String retNum = String.valueOf(whileUseCount);
        whileUseCount++;
        return retNum;
    }

    private String getLastWhileNum() {
        assert whileUseCount >= 0;
        return String.valueOf(whileUseCount - 1);
    }

    public void run() throws IOException {
        String FileName = "target/classes/prog.szc";

        CharStream inputFile = CharStreams.fromFileName(FileName);
        SZCGrammarLexer szcLexer = new SZCGrammarLexer(inputFile);
        CommonTokenStream tokens = new CommonTokenStream(szcLexer);
        SZCGrammarParser szcParser = new SZCGrammarParser(tokens);
        ParseTree tree = szcParser.prog();

        visit(tree);

        System.out.println(outputIR.toString());
    }

    @Override
    public String visitConditionCombine(SZCGrammarParser.ConditionCombineContext ctx) {
        List<ParseTree> childrens = ctx.children;
        if (childrens.size() != 1 && childrens.size() != 3) {
            throw new RuntimeException("visitConditionCombine 孩子个数异常");
        }

        String leftRet = visit(childrens.get(0));
        if (childrens.size() == 1) {
            return leftRet;
        }

        // 短路
        String shortCircuitSym = "@result_" + ((1 << 31) - 1);
        if (!symbolTable.containsSysVar(shortCircuitSym)) {
            String shortCircuitAllocInstr = "\t" + shortCircuitSym + " = alloc i32\n";
            outputIR.append(shortCircuitAllocInstr);
            symbolTable.insertSysVarIntoDeepI(shortCircuitSym);
        }

        String ifNum = getAIfNum();
        String elseStoreInstr = null;
        if (Objects.equals(childrens.get(1).getText(), "且")) {
            String reduceLeftSym = getAUniqueTempSym();
            String reduceLeftInstr = "\t" + reduceLeftSym + " = ne 0, " + leftRet + "\n";
            outputIR.append(reduceLeftInstr);
            String scCondInstr = "\tbr " + reduceLeftSym + ", %then_" + ifNum + ", %else_" + ifNum + "\n";
            outputIR.append(scCondInstr);

            elseStoreInstr = "\tstore 0, " + shortCircuitSym + "\n";
        } else if (Objects.equals(childrens.get(1).getText(), "或")) {
            String reduceLeftSym = getAUniqueTempSym();
            String reduceLeftInstr = "\t" + reduceLeftSym + " = eq 0, " + leftRet + "\n";
            outputIR.append(reduceLeftInstr);
            String scCondInstr = "\tbr " + reduceLeftSym + ", %then_" + ifNum + ", %else_" + ifNum + "\n";
            outputIR.append(scCondInstr);

            elseStoreInstr = "\tstore 1, " + shortCircuitSym + "\n";
        }

        // 执行右半边条件
        outputIR.append("%then_" + ifNum + ":\n");
        String rightRet = visit(childrens.get(2));
        String reduceRightSym = getAUniqueTempSym();
        String reduceRightInstr = "\t" + reduceRightSym + " = ne 0, " + rightRet + "\n";
        outputIR.append(reduceRightInstr);
        String storeRightInstr = "\tstore " + reduceRightSym + ", " + shortCircuitSym + "\n";
        outputIR.append(storeRightInstr);
        String jumpInstr = "\tjump %end_" + ifNum + "\n";
        outputIR.append(jumpInstr);

        String ensureElseInstr = "%else_" + ifNum + ":\n";

        String elseJumpInstr = "\tjump %end_" + ifNum + "\n";
        outputIR.append(ensureElseInstr).append(elseStoreInstr).append(elseJumpInstr);

        String endTagInstr = "%end_" + ifNum + ":\n";
        String retSym = getAUniqueTempSym();
        String endLoadSym = "\t" + retSym + " = load " + shortCircuitSym + "\n";
        outputIR.append(endTagInstr).append(endLoadSym);

        return retSym;
    }

    @Override
    public String visitIfElseStatement(SZCGrammarParser.IfElseStatementContext ctx) {
        List<ParseTree> childrens = ctx.children;
        if (childrens.size() != 5 && childrens.size() != 3) {
            throw new RuntimeException("visitIfElseStatement 孩子个数异常");
        }

        String conditionResult = visit(childrens.get(1));

        String ifNum = getAIfNum();
        String ifInstr = "\tbr " + conditionResult + ", %then_" + ifNum + ", %else_" + ifNum + "\n";
        outputIR.append(ifInstr);
        outputIR.append("%then_" + ifNum + ":\n");
        String ifRet = visit(childrens.get(2));
        if (!Objects.equals(ifRet, "$return")) {
            outputIR.append("\t%jump %end_" + ifNum + "\n");
        }

        outputIR.append("%else_" + ifNum + ":\n");

        if (childrens.size() == 5) {
            String elseRet = visit(childrens.get(4));
            if (!Objects.equals(elseRet, "$return")) {
                outputIR.append("\t%jump %end_" + ifNum + "\n");
            }
        } else if (childrens.size() == 3) {
            outputIR.append("\t%jump %end_" + ifNum + "\n");
        }

        outputIR.append("%end_" + ifNum + ":\n");
        return "";
    }

    @Override
    public String visitProg(SZCGrammarParser.ProgContext ctx) {
        List<ParseTree> childrens = ctx.children;

        for (ParseTree childCtx : childrens) {
            visit(childCtx);
        }
        return "";
    }

    @Override
    public String visitReturnStatement(SZCGrammarParser.ReturnStatementContext ctx) {
        List<ParseTree> childrens = ctx.children;
        String ChildRet = null;
        if (childrens.size() == 1) {
            // todo: return;
            assert false;
        } else if (childrens.size() == 2) {
            ChildRet = visit(childrens.get(1));
        } else assert false;

        outputIR.append("\tret ").append(ChildRet).append("\n");
        return "$return";
    }

    @Override
    public String visitLval(SZCGrammarParser.LvalContext ctx) {
        String varName = ctx.IDENTIFIER().getText();
        String sym = getAUniqueTempSym();
        String loadInstr = "\t" + sym + " = load " + symbolTable.varGetRefFromBlockDeep(varName) + "\n";
        outputIR.append(loadInstr);
        return sym;
    }

    @Override
    public String visitFuncDef(SZCGrammarParser.FuncDefContext ctx) {
        outputIR.append("fun @").append(ctx.funcName().IDENTIFIER()).append("(): i32 {\n");
        outputIR.append("%entry:\n");
        visitBlock(ctx.block());
        outputIR.append("}\n");
        return "";
    }

    @Override
    public String visitBlock(SZCGrammarParser.BlockContext ctx) {
        symbolTable.deepBlockDeep();
        for (SZCGrammarParser.AbstractStatementContext stmtCtx : ctx.abstractStatement()) {
            String visitRet = visitAbstractStatement(stmtCtx);
            if (Objects.equals(visitRet, "$return")) {
                return "$return";
            }
        }
        symbolTable.shallowBlockDeep();
        return "";
    }

    @Override
    public String visitAbstractStatement(SZCGrammarParser.AbstractStatementContext ctx) {
        return visit(ctx.children.get(0));
    }

    @Override
    public String visitPrimaryExp(SZCGrammarParser.PrimaryExpContext ctx) {
        List<ParseTree> childrens = ctx.children;
        if (childrens.size() == 1) {
            String childRet = visit(childrens.get(0));
            String visitRet = childRet;
            if (childRet == null) {
                visitRet = ctx.NUMBER().getText();
            }
            return visitRet;
        } else if (childrens.size() == 3) {
            return visit(childrens.get(1));
        } else {
            throw new RuntimeException("visitPrimaryExp 孩子个数异常");
        }
    }

    @Override
    public String visitUnaryExp(SZCGrammarParser.UnaryExpContext ctx) {
        List<ParseTree> childrens = ctx.children;
        ArrayList<String> childRetValues = new ArrayList<>();
        for (ParseTree childCtx : childrens) {
            String visitRet = visit(childCtx);
            childRetValues.add(visitRet);
        }

        if (childrens.size() == 1) {
            return childRetValues.get(0);
        } else if (childrens.size() == 2) {
            if (Objects.equals(childrens.get(1).getText(), "+")) {
                return visit(childrens.get(1));
            }

            String integInstr = null;
            String retSym = getAUniqueTempSym();
            if (Objects.equals(childrens.get(1).getText(), "-")) {
                integInstr = "\t" + retSym + " = sub 0, " + childRetValues.get(1) + "\n";
            } else {
                // !
                integInstr = "\t" + retSym + " = eq 0, " + childRetValues.get(1) + "\n";
            }

            outputIR.append(integInstr);
            return retSym;
        } else {
            throw new RuntimeException("visitPrimaryExp 孩子个数异常");
        }
    }

    @Override
    public String visitExp(SZCGrammarParser.ExpContext ctx) {
        return visit(ctx.lOrExp());
    }

    @Override
    public String visitAddExp(SZCGrammarParser.AddExpContext ctx) {
        List<ParseTree> childrens = ctx.children;
        ArrayList<String> childRetValues = new ArrayList<>();
        for (ParseTree childCtx : childrens) {
            String visitRet = visit(childCtx);
            childRetValues.add(visitRet);
        }

        if (childrens.size() == 1) {
            return childRetValues.get(0);
        } else if (childrens.size() == 3) {
            String retSym = getAUniqueTempSym();
            String integInstr;
            if (Objects.equals(childrens.get(1).getText(), "+")) {
                integInstr = "\t" + retSym + " = add " + childRetValues.get(0) + ", " + childRetValues.get(2) + "\n";
            } else {
                // -
                integInstr = "\t" + retSym + " = sub " + childRetValues.get(0) + ", " + childRetValues.get(2) + "\n";
            }
            outputIR.append(integInstr);
            return retSym;
        } else {
            throw new RuntimeException("visitAddExp 孩子个数异常");
        }
    }

    @Override
    public String visitMulExp(SZCGrammarParser.MulExpContext ctx) {
        List<ParseTree> childrens = ctx.children;
        ArrayList<String> childRetValues = new ArrayList<>();
        for (ParseTree childCtx : childrens) {
            childRetValues.add(visit(childCtx));
        }

        if (childrens.size() == 1) {
            return childRetValues.get(0);
        } else if (childrens.size() == 3) {
            String retSym = getAUniqueTempSym();
            String integInstr;
            if (Objects.equals(childrens.get(1).getText(), "*")) {
                integInstr = "\t" + retSym + " = mul " + childRetValues.get(0) + ", " + childRetValues.get(2) + "\n";
            } else {
                // /
                integInstr = "\t" + retSym + " = div " + childRetValues.get(0) + ", " + childRetValues.get(2) + "\n";
            }

            outputIR.append(integInstr);
            return retSym;
        } else {
            throw new RuntimeException("visitMulExp 孩子个数异常");
        }
    }

    @Override
    public String visitRelExp(SZCGrammarParser.RelExpContext ctx) {
        List<ParseTree> childrens = ctx.children;
        ArrayList<String> childRetValues = new ArrayList<>();
        for (ParseTree childCtx : childrens) {
            childRetValues.add(visit(childCtx));
        }

        if (childRetValues.size() == 1) {
            return childRetValues.get(0);
        } else if (childRetValues.size() == 3) {
            String retSym = getAUniqueTempSym();
            String integInstr;
            if (Objects.equals(childrens.get(1).getText(), "<")) {
                integInstr = "\t" + retSym + " = lt " + childRetValues.get(0) + ", " + childRetValues.get(2) + "\n";
            } else if (Objects.equals(childrens.get(1).getText(), ">")) {
                integInstr = "\t" + retSym + " = gt " + childRetValues.get(0) + ", " + childRetValues.get(2) + "\n";
            } else if (Objects.equals(childrens.get(1).getText(), ">")) {
                integInstr = "\t" + retSym + " = le " + childRetValues.get(0) + ", " + childRetValues.get(2) + "\n";
            } else if (Objects.equals(childrens.get(1).getText(), ">")) {
                integInstr = "\t" + retSym + " = ge " + childRetValues.get(0) + ", " + childRetValues.get(2) + "\n";
            } else {
                throw new RuntimeException("visitRelExp 符号异常");
            }

            outputIR.append(integInstr);
            return retSym;
        } else {
            throw new RuntimeException("visitRelExp 孩子个数异常");
        }
    }

    @Override
    public String visitEqExp(SZCGrammarParser.EqExpContext ctx) {
        List<ParseTree> childrens = ctx.children;
        ArrayList<String> childRetValues = new ArrayList<>();
        for (ParseTree childCtx : childrens) {
            childRetValues.add(visit(childCtx));
        }

        if (childrens.size() == 1) {
            return childRetValues.get(0);
        } else if (childrens.size() == 3) {
            String retSym = getAUniqueTempSym();
            String integInstr;
            if (Objects.equals(childrens.get(1).getText(), "==")) {
                integInstr = "\t" + retSym + " = eq " + childRetValues.get(0) + ", " + childRetValues.get(2) + "\n";
            } else {
                // !=
                integInstr = "\t" + retSym + " = ne " + childRetValues.get(0) + ", " + childRetValues.get(2) + "\n";
            }
            outputIR.append(integInstr);
            return retSym;
        } else {
            throw new RuntimeException("visitEqExp 孩子个数异常");
        }
    }

    @Override
    public String visitLAndExp(SZCGrammarParser.LAndExpContext ctx) {
        List<ParseTree> childrens = ctx.children;
        ArrayList<String> childRetValues = new ArrayList<>();
        for (ParseTree childCtx : childrens) {
            childRetValues.add(visit(childCtx));
        }

        if (childRetValues.size() == 1) {
            return childRetValues.get(0);
        } else if (childRetValues.size() == 3) {
            String leftSym = getAUniqueTempSym();
            String leftInstr = "\t" + leftSym + " = ne 0, " + childRetValues.get(0) + "\n";

            String rightSym = getAUniqueTempSym();
            String rightInstr = "\t" + rightSym + " = ne 0, " + childRetValues.get(2) + "\n";

            String retSym = getAUniqueTempSym();
            String integInstr = "\t" + retSym + " = and " + leftSym + ", " + rightSym + "\n";

            outputIR.append(leftInstr).append(rightInstr).append(integInstr);

            return retSym;
        } else {
            throw new RuntimeException("visitLAndExp 孩子个数异常");
        }
    }

    @Override
    public String visitLOrExp(SZCGrammarParser.LOrExpContext ctx) {
        List<ParseTree> childrens = ctx.children;
        ArrayList<String> childRetValues = new ArrayList<>();
        for (ParseTree childCtx : childrens) {
            childRetValues.add(visit(childCtx));
        }

        if (childRetValues.size() == 1) {
            return childRetValues.get(0);
        } else if (childRetValues.size() == 3) {
            String leftSym = getAUniqueTempSym();
            String leftInstr = "\t" + leftSym + " = ne 0, " + childRetValues.get(0) + "\n";

            String rightSym = getAUniqueTempSym();
            String rightInstr = "\t" + rightSym + " = ne 0, " + childRetValues.get(2) + "\n";

            String retSym = getAUniqueTempSym();
            String integInstr = "\t" + retSym + " = or " + leftSym + ", " + rightSym + "\n";

            outputIR.append(leftInstr).append(rightInstr).append(integInstr);

            return retSym;
        } else {
            throw new RuntimeException("visitLOrExp 孩子个数异常");
        }
    }

    @Override
    public String visitVarDef(SZCGrammarParser.VarDefContext ctx) {
        String varName = ctx.IDENTIFIER().getText();
        String varRef = symbolTable.varGetRefFromBlockDeep(varName);

        List<ParseTree> childrens = ctx.children;
        String assignmentInstr = null;
        if (childrens.size() == 1) {
            assignmentInstr = "\tstore 0, " + varRef + "\n";
        } else if (childrens.size() == 3) {
            String expRet = visit(childrens.get(2));
            assignmentInstr = "\tstore " + expRet + ", " + varRef + "\n";
        } else assert false;

        outputIR.append(assignmentInstr);
        return "";
    }

    @Override
    public String visitVarDeclStatement(SZCGrammarParser.VarDeclStatementContext ctx) {
        List<ParseTree> childrens = ctx.children;
        String varName = ctx.varDef().IDENTIFIER().getText();
        symbolTable.insertNewVarWithBlockDeep(varName);
        String varRef = symbolTable.varGetRefFromBlockDeep(varName);

        String allocInstr = "\t" + varRef + " = alloc i32\n";
        outputIR.append(allocInstr);

        visit(childrens.get(1));

        return "";
    }

    @Override
    public String visitInitVal(SZCGrammarParser.InitValContext ctx) {
        return visit(ctx.exp());
    }

    @Override
    public String visitWhileStatement(SZCGrammarParser.WhileStatementContext ctx) {
        String whileNum = getAWhileNum();
        String jumpEntryInstr = "\tjump %while_entry_" + whileNum + "\n";
        String entryTagInstr = "%while_entry_" + whileNum + ":\n";
        outputIR.append(jumpEntryInstr).append(entryTagInstr);

        List<ParseTree> childrens = ctx.children;
        String conditionResult = visit(childrens.get(1));
        String ifNum = getAIfNum();
        String ifInstr = "\tbr " + conditionResult + ", %while_body_" + ifNum + ", %end_" + ifNum + "\n";
        outputIR.append(ifInstr);
        outputIR.append("%while_body_" + ifNum + ":\n");
        String ifRet = visit(childrens.get(2));
        if (!Objects.equals(ifRet, "$return")) {
            outputIR.append("\tjump %while_entry_" + ifNum + "\n");
        }

        outputIR.append("%end_" + ifNum + ":\n");
        return "";
    }

    @Override
    public String visitBreakStatement(SZCGrammarParser.BreakStatementContext ctx) {
        String lastWhileNum = getLastWhileNum();
        String jumpInstr = "\tjump %end_" + lastWhileNum + "\n";
        outputIR.append(jumpInstr);
        return "$return";
    }

    @Override
    public String visitContinueStatement(SZCGrammarParser.ContinueStatementContext ctx) {
        String lastWhileNum = getLastWhileNum();
        String jumpInstr = "\tjump %while_body_" + lastWhileNum + "\n";
        outputIR.append(jumpInstr);
        return "$return";
    }
}
